{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "7d5aeac8-97bc-4662-9f7c-8f2e5fdb079e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import re\n",
    "import os\n",
    "import time\n",
    "import jwt\n",
    "import requests\n",
    "\n",
    "def generate_token(apikey: str, exp_seconds: int):\n",
    "    try:\n",
    "        id, secret = apikey.split(\".\")\n",
    "    except Exception as e:\n",
    "        raise Exception(\"invalid apikey\", e)\n",
    "\n",
    "    payload = {\n",
    "        \"api_key\": id,\n",
    "        \"exp\": int(round(time.time() * 1000)) + exp_seconds * 1000,\n",
    "        \"timestamp\": int(round(time.time() * 1000)),\n",
    "    }\n",
    "    return jwt.encode(\n",
    "        payload,\n",
    "        secret,\n",
    "        algorithm=\"HS256\",\n",
    "        headers={\"alg\": \"HS256\", \"sign_type\": \"SIGN\"},\n",
    "    )\n",
    "\n",
    "def get_llm(messages, token):\n",
    "    url = \"https://open.bigmodel.cn/api/paas/v4/chat/completions\"\n",
    "    headers = {\n",
    "      'Content-Type': 'application/json',\n",
    "      'Authorization': token\n",
    "    }\n",
    "\n",
    "    data = {\n",
    "        \"model\": \"glm-4\",\n",
    "        \"messages\": messages\n",
    "    }\n",
    "\n",
    "    response = requests.post(url, headers=headers, json=data)\n",
    "    return response.json()\n",
    "    \n",
    "glm_key = os.getenv('glm_key')\n",
    "glm_key = '7bf001734ef2fd7f7a55bf51dadd7cbb.BMAsoKRDFTmTEPwj'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3fae876-c458-4d6f-9f06-55530d86be57",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# 任务1：初始大模型与Agent"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7e8b5e2-25f9-4915-b2e8-6054acccb84a",
   "metadata": {},
   "source": [
    "略"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dafcb0a2-3629-4758-8f25-ec0451368fde",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "# 任务2：ChatGLM API使用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "f2be523b-62be-4047-90f5-6a0260f81f08",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "def regex_agent(question, token):\n",
    "    prompt_template = '''你是一个专业的python的工程师，擅长编写各种的正则表达式。将下面的要求转换为正则匹配表达式。\n",
    "只能输出正则的pattern字符串，要求能直接使用python的re库执行。\n",
    "不要有其他的输出，包括任何的python代码、说明等。\n",
    "{}\n",
    "    '''.format(question)\n",
    "\n",
    "    return get_llm(prompt_template, token)['choices'][0]['message']['content']\n",
    "# 创建token\n",
    "token = generate_token(glm_key, 5)\n",
    "# 获取正则\n",
    "pattern = regex_agent(\"识别首字母大写单词的正则\", token).replace('`','')\n",
    "# 执行正则\n",
    "text = \"Here is an Example of a Capitalized Word and another Example.\"\n",
    "matches = re.findall(pattern, text)\n",
    "# 打印\n",
    "print(matches)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6118b627-1342-4198-83b0-b22d88daab78",
   "metadata": {},
   "source": [
    "# 任务3：数据库内容解析"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78506dde-a961-4d3c-a671-f6b39e053b45",
   "metadata": {},
   "source": [
    "本次命题的数据有两部分：\n",
    "- 各企业的说明书，以pdf形式存放\n",
    "- 各企业、基金、股票的行情变化，以sqlite存放"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6c16a72-43a7-421f-a1c0-5690fbe2394c",
   "metadata": {},
   "source": [
    "通过分析，sqlite由10张表组成，分别是：基金基本信息、基金份额持有人结构、基金日行情表、基金规模变动表、基金可转债持仓明细、基金债券持仓明细、基金股票持仓明细、港股票日行情表、A股票日行情表、A股公司行业划分表。\n",
    "经过db工具的查询分析，可得出以下数据库关系：\n",
    "![avatar](./img/数据库结构.png)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a8833c5b-26aa-4193-bbc4-c8546efcb978",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\turke\\AppData\\Local\\Temp\\ipykernel_7376\\3176817578.py:4: DeprecationWarning: \n",
      "Pyarrow will become a required dependency of pandas in the next major release of pandas (pandas 3.0),\n",
      "(to allow more performant data types, such as the Arrow string type, and better interoperability with other libraries)\n",
      "but was not found to be installed on your system.\n",
      "If this would cause problems for you,\n",
      "please provide us feedback at https://github.com/pandas-dev/pandas/issues/54466\n",
      "        \n",
      "  import pandas as pd\n",
      "C:\\Users\\turke\\AppData\\Local\\Temp\\ipykernel_7376\\3176817578.py:48: RemovedIn20Warning: Deprecated API features detected! These feature(s) are not compatible with SQLAlchemy 2.0. To prevent incompatible upgrades prior to updating applications, ensure requirements files are pinned to \"sqlalchemy<2.0\". Set environment variable SQLALCHEMY_WARN_20=1 to show all deprecation warnings.  Set environment variable SQLALCHEMY_SILENCE_UBER_WARNING=1 to silence this message. (Background on SQLAlchemy 2.0 at: https://sqlalche.me/e/b8d9)\n",
      "  table_instance = Table(table_name, MetaData(), autoload=True, autoload_with=self.engine)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Table -> A股公司行业划分表\n",
      "Table -> A股票日行情表\n",
      "Table -> 基金份额持有人结构\n",
      "Table -> 基金债券持仓明细\n",
      "Table -> 基金可转债持仓明细\n",
      "Table -> 基金基本信息\n",
      "Table -> 基金日行情表\n",
      "Table -> 基金股票持仓明细\n",
      "Table -> 基金规模变动表\n",
      "Table -> 港股票日行情表\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>name</th>\n",
       "      <th>type</th>\n",
       "      <th>nullable</th>\n",
       "      <th>default</th>\n",
       "      <th>autoincrement</th>\n",
       "      <th>primary_key</th>\n",
       "      <th>distinct</th>\n",
       "      <th>mode</th>\n",
       "      <th>nan_count</th>\n",
       "      <th>max</th>\n",
       "      <th>min</th>\n",
       "      <th>random</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>股票代码</th>\n",
       "      <td>股票代码</td>\n",
       "      <td>TEXT</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>auto</td>\n",
       "      <td>0</td>\n",
       "      <td>5205</td>\n",
       "      <td>990018</td>\n",
       "      <td>0</td>\n",
       "      <td>990018</td>\n",
       "      <td>000001</td>\n",
       "      <td>[002561, 301102, 000713]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>交易日期</th>\n",
       "      <td>交易日期</td>\n",
       "      <td>TEXT</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>auto</td>\n",
       "      <td>0</td>\n",
       "      <td>1096</td>\n",
       "      <td>20211231</td>\n",
       "      <td>0</td>\n",
       "      <td>20211231</td>\n",
       "      <td>20190101</td>\n",
       "      <td>[20190103, 20190120, 20190121]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>行业划分标准</th>\n",
       "      <td>行业划分标准</td>\n",
       "      <td>TEXT</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>auto</td>\n",
       "      <td>0</td>\n",
       "      <td>2</td>\n",
       "      <td>申万行业分类</td>\n",
       "      <td>0</td>\n",
       "      <td>申万行业分类</td>\n",
       "      <td>中信行业分类</td>\n",
       "      <td>[中信行业分类]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>一级行业名称</th>\n",
       "      <td>一级行业名称</td>\n",
       "      <td>TEXT</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>auto</td>\n",
       "      <td>0</td>\n",
       "      <td>45</td>\n",
       "      <td>电子</td>\n",
       "      <td>309830</td>\n",
       "      <td>餐饮旅游</td>\n",
       "      <td>交通运输</td>\n",
       "      <td>[电力设备及新能源, 商贸零售, 建筑]</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>二级行业名称</th>\n",
       "      <td>二级行业名称</td>\n",
       "      <td>TEXT</td>\n",
       "      <td>True</td>\n",
       "      <td>None</td>\n",
       "      <td>auto</td>\n",
       "      <td>0</td>\n",
       "      <td>221</td>\n",
       "      <td>None</td>\n",
       "      <td>1974149</td>\n",
       "      <td>黑色家电Ⅱ</td>\n",
       "      <td>IT服务</td>\n",
       "      <td>[一般零售, 建筑施工Ⅱ, 公交物流]</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "          name  type nullable default autoincrement primary_key distinct  \\\n",
       "股票代码      股票代码  TEXT     True    None          auto           0     5205   \n",
       "交易日期      交易日期  TEXT     True    None          auto           0     1096   \n",
       "行业划分标准  行业划分标准  TEXT     True    None          auto           0        2   \n",
       "一级行业名称  一级行业名称  TEXT     True    None          auto           0       45   \n",
       "二级行业名称  二级行业名称  TEXT     True    None          auto           0      221   \n",
       "\n",
       "            mode nan_count       max       min                          random  \n",
       "股票代码      990018         0    990018    000001        [002561, 301102, 000713]  \n",
       "交易日期    20211231         0  20211231  20190101  [20190103, 20190120, 20190121]  \n",
       "行业划分标准    申万行业分类         0    申万行业分类    中信行业分类                        [中信行业分类]  \n",
       "一级行业名称        电子    309830      餐饮旅游      交通运输            [电力设备及新能源, 商贸零售, 建筑]  \n",
       "二级行业名称      None   1974149     黑色家电Ⅱ      IT服务             [一般零售, 建筑施工Ⅱ, 公交物流]  "
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from typing import Union\n",
    "import traceback\n",
    "from sqlalchemy import create_engine, inspect, func, select, Table, MetaData\n",
    "import pandas as pd\n",
    "\n",
    "class DBParser:\n",
    "    '''DBParser'''\n",
    "    def __init__(self, db_url:str) -> None:\n",
    "        '''初始化\n",
    "        db_url: 数据库链接地址\n",
    "        '''\n",
    "\n",
    "        # 判断数据库类型\n",
    "        if 'sqlite' in db_url:\n",
    "            self.db_type = 'sqlite'\n",
    "        elif 'mysql' in db_url:\n",
    "            self.db_type = 'mysql'\n",
    "\n",
    "        # 链接数据库\n",
    "        self.engine = create_engine(db_url, echo=False)\n",
    "        self.conn = self.engine.connect()\n",
    "        self.db_url = db_url\n",
    "\n",
    "        # 查看表明\n",
    "        self.inspector = inspect(self.engine)\n",
    "        self.table_names = self.inspector.get_table_names()\n",
    "\n",
    "        self._table_fields = {} # 数据表字段\n",
    "        self.foreign_keys = [] # 数据库外键\n",
    "        self._table_sample = {} # 数据表样例\n",
    "\n",
    "        # 依次对每张表的字段进行统计\n",
    "        for table_name in self.table_names:\n",
    "            print(\"Table ->\", table_name)\n",
    "            self._table_fields[table_name] = {}\n",
    "\n",
    "            # 累计外键\n",
    "            self.foreign_keys += [\n",
    "                {\n",
    "                    'constrained_table': table_name,\n",
    "                    'constrained_columns': x['constrained_columns'],\n",
    "                    'referred_table': x['referred_table'],\n",
    "                    'referred_columns': x['referred_columns'],\n",
    "                } for x in self.inspector.get_foreign_keys(table_name)\n",
    "            ]\n",
    "\n",
    "            # 获取当前表的字段信息\n",
    "            table_instance = Table(table_name, MetaData(), autoload=True, autoload_with=self.engine)\n",
    "            table_columns = self.inspector.get_columns(table_name)\n",
    "            self._table_fields[table_name] = {x['name']:x for x in table_columns}\n",
    "\n",
    "            # 对当前字段进行统计\n",
    "            for column_meta in table_columns:\n",
    "                # 获取当前字段\n",
    "                column_instance = getattr(table_instance.columns, column_meta['name'])\n",
    "\n",
    "                # 统计unique\n",
    "                query = select(func.count(func.distinct(column_instance)))\n",
    "                distinct_count = self.conn.execute(query).fetchone()[0]\n",
    "                self._table_fields[table_name][column_meta['name']]['distinct'] = distinct_count\n",
    "\n",
    "                # 统计most frequency value\n",
    "                field_type = self._table_fields[table_name][column_meta['name']]['type']\n",
    "                field_type = str(field_type)\n",
    "                if 'text' in field_type.lower() or 'char' in field_type.lower():\n",
    "                    query = (\n",
    "                        select([column_instance, func.count().label('count')])\n",
    "                        .group_by(column_instance)\n",
    "                        .order_by(func.count().desc())\n",
    "                        .limit(1)\n",
    "                    )\n",
    "                    top1_value = self.conn.execute(query).fetchone()[0]\n",
    "                    self._table_fields[table_name][column_meta['name']]['mode'] = top1_value\n",
    "\n",
    "                # 统计missing个数\n",
    "                query = select(func.count()).filter(column_instance == None)\n",
    "                nan_count = self.conn.execute(query).fetchone()[0]\n",
    "                self._table_fields[table_name][column_meta['name']]['nan_count'] = nan_count\n",
    "\n",
    "                # 统计max\n",
    "                query = select(func.max(column_instance))\n",
    "                max_value = self.conn.execute(query).fetchone()[0]\n",
    "                self._table_fields[table_name][column_meta['name']]['max'] = max_value\n",
    "\n",
    "                # 统计min\n",
    "                query = select(func.min(column_instance))\n",
    "                min_value = self.conn.execute(query).fetchone()[0]\n",
    "                self._table_fields[table_name][column_meta['name']]['min'] = min_value\n",
    "\n",
    "                # 任意取值\n",
    "                query = select(column_instance).limit(10)\n",
    "                random_value = self.conn.execute(query).all()\n",
    "                random_value = [x[0] for x in random_value]\n",
    "                random_value = [str(x) for x in random_value if x is not None]\n",
    "                random_value = list(set(random_value))\n",
    "                self._table_fields[table_name][column_meta['name']]['random'] = random_value[:3]\n",
    "\n",
    "            # 获取表样例（第一行）\n",
    "            query = select([table_instance])\n",
    "            self._table_sample[table_name] = pd.DataFrame([self.conn.execute(query).fetchone()])\n",
    "            self._table_sample[table_name].columns = [x['name'] for x in table_columns]\n",
    "\n",
    "\n",
    "    def get_table_fields(self, table_name) -> pd.DataFrame:\n",
    "        '''获取表字段信息'''\n",
    "        return pd.DataFrame.from_dict(self._table_fields[table_name]).T\n",
    "\n",
    "    def get_data_relations(self) -> pd.DataFrame:\n",
    "        '''获取数据库链接信息（主键和外键）'''\n",
    "        return pd.DataFrame(self.foreign_keys)\n",
    "\n",
    "    def get_table_sample(self, table_name) -> pd.DataFrame:\n",
    "        '''获取数据表样例'''\n",
    "        return self._table_sample[table_name]\n",
    "\n",
    "    def check_sql(self, sql) -> Union[bool, str]:\n",
    "        '''检查sql是否合理\n",
    "\n",
    "        参数\n",
    "            sql: 待执行句子\n",
    "\n",
    "        返回: 是否可以运行 报错信息\n",
    "        '''\n",
    "        try:\n",
    "            self.engine.execute(sql)\n",
    "            return True, 'ok'\n",
    "        except:\n",
    "            err_msg = traceback.format_exc()\n",
    "            return False, err_msg\n",
    "\n",
    "    def execute_sql(self, sql) -> bool:\n",
    "        '''运行SQL'''\n",
    "        result = self.engine.execute(sql)\n",
    "        return list(result)\n",
    "\n",
    "parser = DBParser('sqlite:///./data/博金杯比赛数据.db')\n",
    "parser.get_table_fields(\"A股公司行业划分表\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "eafb62f0-eca0-4420-a089-127e67a04be7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'总共有5205个股票。'"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "token = generate_token(glm_key, 5)\n",
    "# 用户问题\n",
    "question = '查询下总共有多少个股票'\n",
    "'''\n",
    "    第一次询问大模型,获取sql\n",
    "'''\n",
    "prompt_template = '''你是一个sql专家，基于已有的表格信息，请将下面的问题转换为sql查询语句。直接输出sql，不要输出其他内容。\n",
    "表名称：{}\n",
    "\n",
    "表格信息：\n",
    "{}\n",
    "\n",
    "待查询问题：{}\n",
    "'''.format(\"A股公司行业划分表\", parser.get_table_fields(\"A股公司行业划分表\").to_markdown(), question)\n",
    "messages = [{\"role\": \"user\", \"content\": prompt_template}]\n",
    "# 获取sql\n",
    "llm_resp = get_llm(messages, token)['choices'][0]['message']\n",
    "messages.append(llm_resp)\n",
    "sql =llm_resp['content'].replace('`','').replace('sql','')\n",
    "# 执行sql\n",
    "result_list = parser.execute_sql(sql)[0][0]\n",
    "'''\n",
    "    第二次询问大模型，根据数据库结果回答问题。\n",
    "'''\n",
    "# 根据响应回答\n",
    "prompt_template = \"\"\"请根据数据库查询的结果回答我的问题。\n",
    "数据库结果：{}\n",
    "问题：{}\"\"\".format(result_list, question)\n",
    "messages.append({\"role\": \"user\", \"content\": prompt_template})\n",
    "get_llm(messages, token)['choices'][0]['message']['content']"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "546243b4-a4d0-4c92-95a7-3bc2c92f100d",
   "metadata": {},
   "source": [
    "# 任务4：文本索引与答案检索"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53d26ff2-a190-4234-ba45-be58e3f8f34a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
